import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D\

## 从CSV文件加载深度图数据img并显示
img=np.genfromtxt('../../Vision3D_Demos_Data/02.depthImg_to_PC/img_dep_640x480.csv', delimiter=',').astype(np.float32)
print("img.shape: \n", img.shape)
plt.imshow(np.clip(img, 0.55, 0.7),cmap='jet')    # 显示加载的深度图
plt.title('depth image')
plt.show()

## 生成点云使用的相机参数如下：
CAM_WID, CAM_HGT = 640,480           # 深度图img的图像尺寸
CAM_FX, CAM_FY   = 795.209,793.957   # 相机的fx/fy参数
CAM_CX, CAM_CY   = 332.031,231.308   # 相机的cx/cy参数
CAM_DVEC = np.array([-0.33354, 0.00924849, -0.000457208, -0.00215353, 0.0]) # 相机镜头的矫正参数，用于cv2.undistort()的输入之一

"""
index = 0
pc = np.zeros((img.shape[0]*img.shape[1], 3), np.float)
for hIndex in range(0, CAM_HGT):
    for wIndex in range(0, CAM_WID):
        z = img[hIndex, wIndex]
        pc[index, 0] = (wIndex - CAM_CX) * z / CAM_FX; 
        pc[index, 1] = (hIndex - CAM_CY) * z / CAM_FY; 
        pc[index, 2] = z 
        index += 1
"""

# 转换
x, y = np.meshgrid(range(CAM_WID), range(CAM_HGT))
y = y.astype(np.float32) - CAM_CY   # hCoor
x = x.astype(np.float32) - CAM_CX   # wCoor

img_z = img.copy()
if False: # 如果需要矫正视线到Z的转换的话使能
    f = (CAM_FX + CAM_FY)/2.0
    img_z *= f/np.sqrt(x**2 + y**2 + f**2)

# 以矩阵的方式计算x,y坐标 
pc_x = img_z * x/CAM_FX # X = Z * (u - cx)/fx 
pc_y = img_z * y/CAM_FY # Y = Z * (v - cy)/fy 
pc = np.array([pc_x.ravel(), pc_y.ravel(), img_z.ravel()]).T
print("pc.shape: ", pc.shape)

np.savetxt('../../Vision3D_Demos_Data/02.depthImg_to_PC/pc.csv', pc, fmt='%.18e', delimiter=',', newline='\n')

# 下面是保存CSV代码的例子以及显示点云的例子
# 例子：生成并保存随机点云，保存为CSV文件
# pc=np.random.rand(5000,3)*2.0-1.0 # 生成随机点云
# np.savetxt('pc.csv', pc, fmt='%.18e', delimiter=',', newline='\n')

## 从CSV文件加载点云并显示
pc=np.genfromtxt('../../Vision3D_Demos_Data/02.depthImg_to_PC/pc.csv', delimiter=',').astype(np.float32)
ax = plt.figure(1).gca(projection='3d')
ax.plot(pc[:,0], pc[:,1], pc[:,2],'b.', markersize=0.5)
plt.title('point cloud')
plt.show()    

